
from rag import load_pdf
from base import Agent
import streamlit as st
import requests

future_work = '''- Role: 科研文献探索点分析师
- Background: 用户需要从科研文献中识别出可以进一步探索的点，以促进未来的研究工作或拓展当前研究的深度和广度。
- Profile: 你是一位专业的科研文献探索点分析师，拥有丰富的科研背景和文献分析经验，能够准确识别文献中提出的新问题、未解决的问题以及潜在的研究领域。
- Skills: 你具备出色的文献解读能力、批判性思维和前瞻性分析能力，能够从论文中提取出未来研究的潜在方向和关键问题。
- Goals: 从提供的科研论文中提取出可以进一步探索的点，包括但不限于新的研究方向、未解决的问题、数据集的扩展、方法论的改进等。
- Constrains: 提取的探索点应基于文献内容，避免主观臆断，确保提出的点具有实际的研究价值和可行性。
- OutputFormat: 列出可以进一步探索的点，并为每个点提供简短的描述和可能的研究建议。
- Workflow:
  1. 仔细阅读科研论文，特别是讨论和结论部分。
  2. 识别文献中提出的新问题、未解决的问题以及作者建议的未来研究方向。
  3. 分析文献中提及的数据集、方法论和实验设计的局限性。
  4. 根据文献内容，提出可能的改进措施和新的研究点。
  5. 清晰、准确地列出可以进一步探索的点，并提供简短描述。
文章内容：
'''

def download_pdf(arxiv_id):
    # download the pdf file
    url = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
    response = requests.get(url)
    with open(f"download/{arxiv_id}.pdf", "wb") as f:
        f.write(response.content)

# load pdf file

def read_pdf(arxiv_id):
    docs = load_pdf(f"download/{arxiv_id}.pdf")
    
    pages = []
    for page in docs:
        pages.append(page.page_content)

    return " ".join(pages)

# create vector database
def get_future_work(arxiv_id):
    pdf_content = read_pdf(arxiv_id)
    agent = Agent(future_work)
    return agent.chat("文章内容:"+pdf_content+"\n提取信息：", stream=True)



if __name__ == "__main__":
    st.title("科研问题发现")
    input_question = st.text_area("请输入您想分析的论文的arxiv id:", height=200)

    if st.button("提交"):
        if input_question:
            arxiv_id = input_question
            download_pdf(arxiv_id)
            future_work = get_future_work(arxiv_id)
            st.write_stream(future_work)